package cbbx_sm.decision_maker.search;

import java.util.HashSet;

public class TreeLookAhead {

	public static void growTree(SearchNode root, WorldStateTransition worldStateTransition,int levels) {
		SearchDriver driver = new SearchDriver(root);
		
		HashSet<SearchNode> frontierNodes = null;
		//System.out.println(" Num States = "+(Schedules.getSchedules().length-1));
		for (int i=1; i<=levels; i++){
			frontierNodes = driver.expandFrontierNodes(false);
		}
		
		/** Perform a backwards walk which takes into account only actions which are rational by the optimal policy and compute the maximal score. */
		/** For the leaf nodes, set the score as we already know what it is. */
		for (SearchNode node: frontierNodes){
			ScheduleConfigurationNode scnode = (ScheduleConfigurationNode)node;
			scnode.setScore(0);
		}
		
		while (frontierNodes.size()!=0){
			HashSet<SearchNode> currentLevel = new HashSet<SearchNode>();
			
			for (SearchNode node: frontierNodes){
				ScheduleConfigurationNode parent = ((ScheduleConfigurationNode)node).getParent();
				if (parent!=null){
					currentLevel.add(parent);
				}
			}
			
			/** For each parent compute the action that would be taken based on its policy. 
			 *  Add to the score of that parent the expected utility of that action based on the transition probabilities and expected utilities of parents. */
			for (SearchNode node: currentLevel){
				ScheduleConfigurationNode scnode = (ScheduleConfigurationNode)node;
				double[] scores = new double[worldStateTransition.getSchedules().length];
				for (SearchNode child: scnode.getSuccessors()){
					ScheduleConfigurationNode schChild = (ScheduleConfigurationNode)child;
					scores[worldStateTransition.scheduleToCameraId(schChild.getSelectedAction())] +=
						schChild.getTransitionProbability()*(schChild.getScore()+schChild.getUtility());
				}
				
				double maxScore = 0;
				for (int i=0; i<scores.length; i++){
					maxScore = Math.max(maxScore, scores[i]);
					if (maxScore == scores[i]){
						scnode.setScore(maxScore);
						scnode.setBestNextAction(worldStateTransition.getSchedules()[i]);
					}
				}
			}
			
			frontierNodes = (HashSet<SearchNode>)currentLevel.clone();
		}	
	}
}

